
import torch,shutil,os
from tensorboardX import SummaryWriter
# 计算top值
def get_accuracy(output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))

            res = []
            for k in topk:
                    correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                    res.append(correct_k.mul_(100.0 / batch_size))
            return res
            
def save_checkpoint(state, is_best, model_name='model', path='checkpoints'):
    torch.save(state, os.path.join(path, str(state['epoch'])+'_'+model_name+'_'+str(round(state['acc1'],2))+'.pth'))
    if is_best:
        torch.save(state, os.path.join(path, 'model_best.pth'))
        
# 更改学习率
def adjust_learning_rate(optimizer, epoch, lr, setlr):
    """Sets """
    if epoch in setlr:
        for param_group in optimizer.param_groups:
            lr_data = lr[setlr.index(epoch) + 1]
            print("epoch ", epoch,"set lr is ", lr_data)
            param_group['lr'] = lr[setlr.index(epoch) + 1]
class AverageMeter(object):
    """计算并存储平均值和当前值
       Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262
    """
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count 
class logger(object):
    '''
    进行日志的输出
    '''
    def __init__(self, logdir="logs"):
        self.writer = SummaryWriter(log_dir=logdir)
        
    def add_model(self, model, input_data):
        self.writer.add_graph(model, (input_data,))
        
    def add_data(self, name, data, epoch):
        self.writer.add_scalar(name, data, epoch)
        
    def add_train_data(self, data, epoch):
        self.add_data('loss/loss', data[0], epoch)
        self.add_data('acc/acc1', data[1], epoch)
        self.add_data('acc/acc5', data[2], epoch)
        
    def add_train_in_data(self, loss, epoch):
        self.add_data('loss/stop_loss', loss, epoch)
        
    def close(self):
        self.writer.close()